from manim import *

class ArithmeticSeriesSquares(Scene):
    def construct(self):
        # 设置等差数列的第一个元素和公差
        a1 = 1
        d = 1
        n = 6  # 项数

        # 计算最大值以动态调整单位长度，确保坐标轴在屏幕内
        max_value = a1 + (n - 1) * d
        unit_length = 3 / max_value  # 3 是一个经验值，可以根据需要调整

        # 创建坐标轴，扩大范围
        axes = Axes(
            x_range=[0, n + 2, 1],  # x 轴范围稍微扩大一些
            y_range=[0, max_value + 2, 1],  # y 轴范围稍微扩大一些
            x_length=(n + 2) * unit_length,
            y_length=(max_value + 2) * unit_length,
            axis_config={"include_numbers": True},
        ).add_coordinates().shift(UP * 0.5)
        axes_labels = axes.get_axis_labels(x_label="n", y_label="a_n")
        self.play(Create(axes), Write(axes_labels))

        # 添加等差数列的通项公式
        tex_template = TexTemplate()
        tex_template.add_to_preamble(r"\usepackage{ctex}")
        general_term = MathTex(
            r"a_n = a_1 + (n-1) \cdot d",
            tex_template=tex_template
        ).to_edge(UP)
        self.play(Write(general_term))
        self.wait(0.5)
        self.play(FadeOut(general_term))

        # 添加首项的说明
        first_term = Tex(f"首项 $a_1 = {a1}$ 公差 $d = {d}$ 项数 $n = {n}$", color=WHITE, tex_template=tex_template).to_edge(UP)
        self.play(Write(first_term))
        self.wait(1)
        self.play(FadeOut(first_term))

        # 创建等差数列的正方形表示
        squares = VGroup()
        for i in range(n):
            num_squares = a1 + i * d
            # 提示当前数列项的值
            current_term = Tex(f"$a_{i+1} = {num_squares}$", color=WHITE, tex_template=tex_template).to_edge(UP)
            self.play(Write(current_term))
            for j in range(num_squares):
                square = Square(side_length=unit_length, color=BLUE, fill_opacity=0.5)
                square.move_to(axes.c2p(i + 0.5, j + 0.5))
                squares.add(square)
            self.play(Create(squares[-num_squares:]))  # 只显示当前项的正方形
            self.wait(0.2)
            self.play(FadeOut(current_term))
        
        # 添加说明文字
        explanation1 = Tex("数列的和就是所有正方形的面积和", color=WHITE, tex_template=tex_template).to_edge(UP)
        self.play(Write(explanation1))
        self.wait(2)
        self.play(FadeOut(explanation1))
        
        # 添加说明文字
        explanation2 = Tex("复制所有正方形，刚好可以拼成一个长方形", color=WHITE, tex_template=tex_template).to_edge(UP)
        self.play(Write(explanation2))
        self.wait(2)
        self.play(FadeOut(explanation2))        

        # 复制正方形并旋转180度
        copied_squares = squares.copy()
        copied_squares.rotate(PI, about_point=axes.c2p(n / 2, max_value / 2))
        copied_squares.shift(UP * unit_length)
        copied_squares.set_color(RED)

        # 逐项突出显示原来的正方形并显示复制的正方形
        index = 0
        for i in range(n):
            num_squares = a1 + i * d
            original_squares = squares[index:index + num_squares]
            copied_squares_part = copied_squares[index:index + num_squares]
            
            # 突出显示原来的正方形
            self.play(Indicate(original_squares, scale_factor=1.2, color=YELLOW), run_time=1)
            
            # 显示复制的正方形
            self.play(Create(copied_squares_part), run_time=1)
            
            index += num_squares

        # 添加说明文字
        explanation3 = Tex("长方形的面积刚好是数列和的2倍", color=WHITE, tex_template=tex_template).to_edge(UP)
        self.play(Write(explanation3))
        self.wait(2)
        self.play(FadeOut(explanation3))
        
        # 突出显示矩形的宽和高
        width_label = Tex(f"宽 = n = {n}", color=YELLOW, tex_template=tex_template).next_to(axes.c2p(n / 2, -1), DOWN)
        height_label = Tex(f"高 = $a_1 + a_n = {a1 + max_value}$", color=YELLOW, tex_template=tex_template).next_to(axes.c2p(-1, max_value / 2), LEFT)
        width_line = Line(axes.c2p(0, -0.5), axes.c2p(n, -0.5), color=YELLOW)
        height_line = Line(axes.c2p(-0.5, 0), axes.c2p(-0.5, max_value+1), color=YELLOW)

        self.play(Create(width_line), Write(width_label))
        self.play(Create(height_line), Write(height_label))
        self.wait(2)
        self.play(FadeOut(width_label),FadeOut(height_label),FadeOut(width_line),FadeOut(height_line))
        
        # 添加求和公式
        sum_formula = MathTex(
            r"S_n = \frac{n \times (a_1 + a_n) }{2}",
            tex_template=tex_template
        ).to_edge(DOWN)
        self.play(Write(sum_formula))
        
        # 保持动画一段时间
        self.wait(2)

# 运行代码
# 保存修改后的 Python 文件（例如 ArithmeticSeriesSquares.py），然后在命令行中运行以下命令：
# manim -pql .\ArithmeticSeriesSquares.py